from fastapi import FastAPI, File, UploadFile, Form, Query
from fastapi.responses import JSONResponse, FileResponse
from pathlib import Path
import shutil
import os
from diffusers import StableDiffusionPipeline
import torch

app = FastAPI()

# 上传目录
UPLOAD_DIR = Path("training_data")
UPLOAD_DIR.mkdir(parents=True, exist_ok=True)

# 文生图输出目录
OUTPUT_DIR = Path("generated")
OUTPUT_DIR.mkdir(parents=True, exist_ok=True)

# 加载模型（首次加载较慢）
pipe = StableDiffusionPipeline.from_pretrained(
    "runwayml/stable-diffusion-v1-5",
    torch_dtype=torch.float16 if torch.backends.mps.is_available() else torch.float32
)
pipe = pipe.to("mps" if torch.backends.mps.is_available() else "cpu")

@app.post("/upload")
async def upload_image(label: str = Form(...), image: UploadFile = File(...)):
    # 创建标签目录
    label_dir = UPLOAD_DIR / label
    label_dir.mkdir(parents=True, exist_ok=True)

    # 保存图像
    image_path = label_dir / image.filename
    with open(image_path, "wb") as buffer:
        shutil.copyfileobj(image.file, buffer)

    return JSONResponse({
        "message": "图像上传成功",
        "label": label,
        "filename": image.filename
    })

@app.get("/generate")
def generate_image(prompt: str = Query(...)):
    # 生成图像
    image = pipe(prompt).images[0]

    # 保存图像
    output_path = OUTPUT_DIR / "output.png"
    image.save(output_path)

    return FileResponse(output_path, media_type="image/png")

@app.get("/")
def root():
    return {"message": "Scaffold AI 接口运行中：支持上传和图像生成"}
